"""
data_manager.py
===============

📦 DataManager 模块

功能：
1. 自动加载数据集
   - 支持内置数据集（如 CIFAR10）
   - 支持本地自定义数据集
   - 根据配置文件自动判断数据是否已下载，未下载则自动下载

2. 构建 DataLoader
   - 根据 DatasetConfig 自动构建训练集、验证集、测试集的 DataLoader
   - 支持批大小、线程数、是否打乱数据、是否固定内存等参数

3. 数据拆分与预处理
   - 内置数据集可按比例拆分 train/val/test
   - 支持数据标准化（normalize）
   - 支持图片缩放（resize）
   - 可选择缓存数据以加速训练

使用方法示例：
----------------
from dataset_config import DatasetConfig
from data_manager import DataManager

# 配置数据集
config = DatasetConfig(name="CIFAR10", data_modal="image", batch_size=32)

# 构建 DataManager
dm = DataManager(config)

# 查看 DataLoader 信息
dm.summary()

# 遍历训练集
for images, labels in dm.train_loader:
    print(images.shape, labels.shape)
    break

注意事项：
----------------
1. dataset_type: "builtin" 或 "local"
2. 内置数据集没有独立验证集时使用 split_ratio 拆分
3. 本地数据集路径需存在且按类别组织子文件夹
4. 对于大数据集，可开启 cache_data；num_workers 根据 CPU 核心数设置
"""


from torch.utils.data import DataLoader

from config_manager.dataset_config import DatasetConfig
from .datasets_manager import DatasetManager


class DataManager:
    def __init__(self, dataset_config:DatasetConfig):
        self.dataset_config = dataset_config
        self.dataset_manager = DatasetManager(self.dataset_config)

    def _make_loader(self, ds, shuffle):
        return DataLoader(
            ds, batch_size=self.dataset_config.batch_size, shuffle=shuffle,
            num_workers=self.dataset_config.num_workers, pin_memory=self.dataset_config.pin_memory,
            drop_last=self.dataset_config.drop_last
        ) if ds is not None else None

    def get_dataloader(self):

        train_loader=self._make_loader(self.dataset_manager.dataset["train"], self.dataset_config.shuffle)
        val_loader=self._make_loader(self.dataset_manager.dataset["val"], False)
        test_loader=self._make_loader(self.dataset_manager.dataset["test"], False)
        return train_loader,val_loader,test_loader





